import itertools
import math
import os
import typing
from dataclasses import dataclass

import hydra.utils
import lightning as L
import numpy as np
import torch
import torch.nn.functional as F
import torchmetrics
import transformers
from torch import Tensor

import dataloader
import models
import noise_schedule
import utils
from framework.hybrid_tran.models.vqgan import Chroma_VQGAN

LOG2 = math.log(2)


def _sample_categorical(categorical_probs):
  gumbel_norm = (
    1e-10
    - (torch.rand_like(categorical_probs) + 1e-10).log())
  return (categorical_probs / gumbel_norm).argmax(dim=-1)


def _unsqueeze(x, reference):
  return x.view(
    * x.shape,
    * ((1,) * (len(reference.shape) - len(x.shape))))


@dataclass
class Loss:
  loss: torch.FloatTensor
  nlls: torch.FloatTensor
  token_mask: torch.FloatTensor


class NLL(torchmetrics.aggregation.MeanMetric):
  pass


class BPD(NLL):
  def compute(self) -> Tensor:
    """Computes the bits per dimension.

    Returns:
      bpd
    """
    return self.mean_value / self.weight / LOG2


class Perplexity(NLL):
  def compute(self) -> Tensor:
    """Computes the Perplexity.

    Returns:
     Perplexity
    """
    return torch.exp(self.mean_value / self.weight)


class Diffusion(L.LightningModule):
  def __init__(
    self,
    config,
    tokenizer: transformers.PreTrainedTokenizer,
    sample_mode = False):
    super().__init__()
    self.save_hyperparameters()
    self.config = config

    self.tokenizer = tokenizer
    # self.vocab_size = self.tokenizer.vocab_size
    self.vocab_size = 16384
    if sample_mode:
      self.chroma_vqgan = Chroma_VQGAN(model_config=None, model_path=None).eval().requires_grad_(False)
    else:
      self.chroma_vqgan = None
    self.sampler = self.config.sampling.predictor
    if not sample_mode:
      self.gen_ppl_eval_model_name_or_path = self.config.eval.\
        gen_ppl_eval_model_name_or_path
    self.antithetic_sampling = self.config.training.antithetic_sampling
    self.importance_sampling = self.config.training.importance_sampling
    self.change_of_variables = self.config.training.change_of_variables
    if (not hasattr(self.tokenizer, 'mask_token')
        or self.tokenizer.mask_token is None):
      self.mask_index = self.vocab_size
      self.vocab_size += 1
    else:
      self.mask_index = self.tokenizer.mask_token_id
    self.parameterization = self.config.parameterization
    if self.config.backbone == 'dit':
      self.backbone = models.dit.DIT(
        self.config, vocab_size=self.vocab_size)
    elif self.config.backbone == 'dimamba':
      self.backbone = models.dimamba.DiMamba(
        self.config,
        vocab_size=self.vocab_size,
        pad_token_id=self.tokenizer.pad_token_id)
    elif self.config.backbone == 'ar':
      self.backbone = models.autoregressive.AR(
        self.config,
        vocab_size=self.vocab_size,
        mask_index=self.mask_index)
    elif self.config.backbone == 'hf_dit':
      self.backbone = transformers.AutoModelForMaskedLM.from_pretrained(
        config.eval.checkpoint_path, trust_remote_code=True)
    else:
      raise ValueError(
        f'Unknown backbone: {self.config.backbone}')

    self.T = self.config.T
    self.subs_masking = self.config.subs_masking

    self.softplus = torch.nn.Softplus()
    # metrics are automatically reset at end of epoch
    metrics = torchmetrics.MetricCollection({
      'nll': NLL(),
      'bpd': BPD(),
      'ppl': Perplexity(),
    })
    metrics.set_dtype(torch.float64)
    self.train_metrics = metrics.clone(prefix='train/')
    self.valid_metrics = metrics.clone(prefix='val/')
    self.test_metrics = metrics.clone(prefix='test/')
    if not sample_mode:
      # generative perplexity
      self.gen_ppl_metric = Perplexity()
      self.eval_model_tokenizer = transformers.AutoTokenizer.\
        from_pretrained(self.gen_ppl_eval_model_name_or_path)
      if self.eval_model_tokenizer.pad_token is None:
        self.eval_model_tokenizer.pad_token =\
            self.eval_model_tokenizer.eos_token
        self.eval_model_tokenizer.pad_token_id =\
            self.eval_model_tokenizer.eos_token_id

    self.noise = noise_schedule.get_noise(self.config,
                                          dtype=self.dtype)
    if self.config.training.ema > 0:
      self.ema = models.ema.ExponentialMovingAverage(
        itertools.chain(self.backbone.parameters(),
                        self.noise.parameters()),
        decay=self.config.training.ema)
    else:
      self.ema = None
    if not sample_mode:
      self.lr = self.config.optim.lr
    self.sampling_eps = self.config.training.sampling_eps
    self.time_conditioning = self.config.time_conditioning
    self.neg_infinity = -1000000.0
    self.fast_forward_epochs = None
    self.fast_forward_batches = None
    self._validate_configuration()

  def _validate_configuration(self):
    assert not (self.change_of_variables
                and self.importance_sampling)
    if self.parameterization == 'sedd':
      assert not self.importance_sampling
      assert not self.change_of_variables
    if self.parameterization == 'd3pm':
      assert self.T > 0
    if self.T > 0:
      assert self.parameterization in {'d3pm', 'subs'}
    if self.subs_masking:
      assert self.parameterization == 'd3pm'

  def _log_visualization_to_wandb(self, colored_img_norm, noisy_img_norm, output_img_norm,
                                 color_cond, cond_mask, cond_tokens, color_condition_save, f_gray):
    """上传可视化图像到WandB"""
    if not (hasattr(self.trainer, 'logger') and self.trainer.logger is not None):
      return
    
    try:
      import wandb
      import torchvision.transforms as T
      
      # 只上传前2张图像避免过多
      for i in range(min(4, colored_img_norm.shape[0])):
        # 构建组合图像
        if color_cond is not None and cond_mask is not None and cond_mask.any():
          if cond_tokens is not None:
            combined = torch.cat([colored_img_norm[i], noisy_img_norm[i], 
                               color_condition_save[i], output_img_norm[i]], dim=2)
          else:
            # 需要重新计算cond_img_norm
            cond_img = self.chroma_vqgan.decode_img(color_cond, f_gray)
            cond_img_norm = torch.clamp((cond_img + 1) / 2, 0, 1)
            combined = torch.cat([colored_img_norm[i], noisy_img_norm[i], 
                               cond_img_norm[i], output_img_norm[i]], dim=2)
        else:
          combined = torch.cat([colored_img_norm[i], noisy_img_norm[i], output_img_norm[i]], dim=2)
        
        # 转换为PIL图像
        pil_image = T.ToPILImage()(combined)
        
        # 上传到WandB
        self.trainer.logger.experiment.log({
          f"visualization/sample_{i}": wandb.Image(
            pil_image, 
            caption=f"Step {self.global_step} - Sample {i} (Original|Noisy|Condition|Output)"
          )
        }, step=self.global_step)
        
    except Exception as e:
      print(f"Failed to upload visualization to WandB: {e}")

  def setup(self, stage=None):
    """Setup hook called by PyTorch Lightning before any stage (fit, validate, test, predict).
    This ensures the VQGAN is initialized properly on the correct device in DDP settings."""
    if self.chroma_vqgan is None:
      self.chroma_vqgan = Chroma_VQGAN(model_config=None, model_path=None).eval().requires_grad_(False)
      # Move to the same device as the current module
      if self.device:  # device may not be set yet during initialization
        self.chroma_vqgan = self.chroma_vqgan.to(self.device)

  def on_load_checkpoint(self, checkpoint):
    if self.ema:
      self.ema.load_state_dict(checkpoint['ema'])
    # Copied from:
    # https://github.com/Dao-AILab/flash-attention/blob/main/training/src/datamodules/language_modeling_hf.py#L41
    self.fast_forward_epochs = checkpoint['loops'][
      'fit_loop']['epoch_progress']['current']['completed']
    self.fast_forward_batches = checkpoint['loops'][
      'fit_loop']['epoch_loop.batch_progress'][
        'current']['completed']

  def on_save_checkpoint(self, checkpoint):
    if self.ema:
      checkpoint['ema'] = self.ema.state_dict()
    # Copied from:
    # https://github.com/Dao-AILab/flash-attention/blob/main/training/src/tasks/seq.py
    # ['epoch_loop.batch_progress']['total']['completed'] is 1 iteration
    # behind, so we're using the optimizer's progress.
    checkpoint['loops']['fit_loop'][
      'epoch_loop.batch_progress']['total'][
        'completed'] = checkpoint['loops']['fit_loop'][
          'epoch_loop.automatic_optimization.optim_progress'][
            'optimizer']['step']['total'][
              'completed'] * self.trainer.accumulate_grad_batches
    checkpoint['loops']['fit_loop'][
      'epoch_loop.batch_progress']['current'][
        'completed'] = checkpoint['loops']['fit_loop'][
          'epoch_loop.automatic_optimization.optim_progress'][
            'optimizer']['step']['current'][
              'completed'] * self.trainer.accumulate_grad_batches
    # _batches_that_stepped tracks the number of global steps, not the number
    # of local steps, so we don't multiply with self.trainer.accumulate_grad_batches here.
    checkpoint['loops']['fit_loop'][
      'epoch_loop.state_dict'][
        '_batches_that_stepped'] = checkpoint['loops']['fit_loop'][
          'epoch_loop.automatic_optimization.optim_progress'][
            'optimizer']['step']['total']['completed']
    if 'sampler' not in checkpoint.keys():
      checkpoint['sampler'] = {}
    if hasattr(self.trainer.train_dataloader.sampler,
               'state_dict'):
      sampler_state_dict = self.trainer.\
        train_dataloader.sampler.state_dict()
      checkpoint['sampler'][
        'random_state'] = sampler_state_dict.get(
          'random_state', None)
    else:
      checkpoint['sampler']['random_state'] = None

  def on_train_start(self):
    if self.ema:
      self.ema.move_shadow_params_to_device(self.device)
    # Make sure VQGAN is on the correct device after the device is determined
    if self.chroma_vqgan is not None and hasattr(self, 'device'):
      self.chroma_vqgan = self.chroma_vqgan.to(self.device)
    # Adapted from:
    # https://github.com/Dao-AILab/flash-attention/blob/main/training/src/datamodules/language_modeling_hf.py
    distributed = (
      self.trainer._accelerator_connector.use_distributed_sampler
      and self.trainer._accelerator_connector.is_distributed)
    if distributed:
      sampler_cls = dataloader.FaultTolerantDistributedSampler
    else:
      sampler_cls = dataloader.RandomFaultTolerantSampler
    updated_dls = []
    for dl in self.trainer.fit_loop._combined_loader.flattened:
      if hasattr(dl.sampler, 'shuffle'):
        dl_sampler = sampler_cls(
          dl.dataset, shuffle=dl.sampler.shuffle)
      else:
        dl_sampler = sampler_cls(dl.dataset)
      if (distributed
          and self.fast_forward_epochs is not None
          and self.fast_forward_batches is not None):
        dl_sampler.load_state_dict({
          'epoch': self.fast_forward_epochs,
          'counter': (self.fast_forward_batches
                      * self.config.loader.batch_size)})
      updated_dls.append(
        torch.utils.data.DataLoader(
          dl.dataset,
          batch_size=self.config.loader.batch_size,
          num_workers=self.config.loader.num_workers,
          pin_memory=self.config.loader.pin_memory,
          sampler=dl_sampler,
          shuffle=False,
          persistent_workers=True))
    self.trainer.fit_loop._combined_loader.flattened = updated_dls

  def optimizer_step(self, *args, **kwargs):
    super().optimizer_step(*args, **kwargs)
    if self.ema:
      self.ema.update(itertools.chain(
        self.backbone.parameters(),
        self.noise.parameters()))

  def _subs_parameterization(self, logits, xt, cond_mask=None, color_cond=None):
    # log prob at the mask index = - infinity
    logits[:, :, self.mask_index] += self.neg_infinity
    
    # Normalize the logits such that x.exp() is
    # a probability distribution over vocab_size.
    logits = logits - torch.logsumexp(logits, dim=-1,
                                      keepdim=True)

    # Apply updates directly in the logits matrix.
    # For the logits of the unmasked tokens, set all values
    # to -infinity except for the indices corresponding to
    # the unmasked tokens.
    unmasked_indices = (xt != self.mask_index)
    
    # 处理条件掩码区域 - 这些区域需要计算loss
    if cond_mask is not None and color_cond is not None:
      # 创建条件区域掩码
      cond_positions = (cond_mask == 1).reshape(cond_mask.shape[0], -1)
      
      # 对于非条件区域的未掩码token，将logits设为-infinity
      # 然后将对应token的位置设为0 (标准处理)
      non_cond_unmasked = unmasked_indices & (~cond_positions)
      logits[non_cond_unmasked] = self.neg_infinity
      logits[non_cond_unmasked, xt[non_cond_unmasked]] = 0
      
      # 对于条件区域：保留原始logits以计算loss
      # 这样模型会学习到在条件区域预测正确的颜色
      # 可以选择性地增强条件区域的权重
      if torch.any(cond_positions & unmasked_indices):
        cond_weight = 1.0  # 可以增加这个权重以强调条件区域的重要性
        cond_unmasked = unmasked_indices & cond_positions
        # 保留原始logits，但我们仍然希望在正确的token上有较大的概率
        # 对于正确的token，给一个较小的boost
        logits[cond_unmasked, xt[cond_unmasked]] += cond_weight
    else:
      # 原始处理方式（无条件）
      logits[unmasked_indices] = self.neg_infinity
      logits[unmasked_indices, xt[unmasked_indices]] = 0
      
    return logits

  def _d3pm_parameterization(self, logits, cond_mask=None, color_cond=None):
    if self.subs_masking:
      logits[:, :, self.mask_index] += self.neg_infinity
    logits = logits - torch.logsumexp(logits, dim=-1,
                                      keepdim=True)
    
    # 如果有条件掩码和颜色条件，对它们进行处理
    if cond_mask is not None and color_cond is not None:
      cond_positions = (cond_mask == 1)
      if torch.any(cond_positions):
        # 对于条件区域，我们希望增强正确颜色的概率
        cond_weight = 1.0  # 可以调整这个权重
        for b in range(logits.shape[0]):
          for pos in torch.nonzero(cond_positions[b]):
            pos_idx = pos.item()
            if pos_idx < color_cond.shape[1]:  # 确保pos_idx在color_cond范围内
              color_idx = color_cond[b, pos_idx].item()
              # 给正确颜色一个boost
              logits[b, pos_idx, color_idx] += cond_weight
              
    return logits

  def _sedd_parameterization(self, logits, xt, sigma, cond_mask=None, color_cond=None):
    esigm1_log = torch.where(
      sigma < 0.5,
      torch.expm1(sigma),
      sigma.exp() - 1).log().to(logits.dtype)
    # logits shape
    # (batch_size, diffusion_model_input_length, vocab_size)
    logits = logits - esigm1_log[:, None, None] - np.log(
      logits.shape[-1] - 1)
    # The below scatter operation sets the log score
    # for the input word to 0.
    logits = torch.scatter(logits, -1, xt[..., None],
                           torch.zeros_like(logits[..., :1]))
    
    # 处理条件掩码区域
    if cond_mask is not None and color_cond is not None:
      cond_positions = (cond_mask == 1)
      if torch.any(cond_positions):
        # 对条件区域增强权重
        cond_weight = 1.0  # 可调整的权重
        for b in range(logits.shape[0]):
          for pos in torch.nonzero(cond_positions[b]):
            pos_idx = pos.item()
            if pos_idx < color_cond.shape[1]:  # 确保pos_idx在color_cond范围内
              color_idx = color_cond[b, pos_idx].item()
              # 对指定的颜色条件增强权重
              if xt[b, pos_idx] == color_idx:
                # 正确的颜色索引，给一个额外的奖励
                logits[b, pos_idx, color_idx] += cond_weight
    
    return logits

  def _process_sigma(self, sigma):
    if sigma is None:
      assert self.parameterization == 'ar'
      return sigma
    if sigma.ndim > 1:
      sigma = sigma.squeeze(-1)
    if not self.time_conditioning:
      sigma = torch.zeros_like(sigma)
    assert sigma.ndim == 1, sigma.shape
    return sigma

  def forward(self, x, sigma, f_gray, x0=None, color_cond=None, cond_mask=None, color_condition=None):
    """Returns log score.
    Args:
      x: Token indices representing current state
      sigma: Noise level 
      f_gray: Grayscale image features
      x0: Original color tokens (for visualization)
      color_cond: Color condition tokens to guide generation
      cond_mask: Mask indicating where color conditions should be applied
    """
    sigma = self._process_sigma(sigma)
    with torch.cuda.amp.autocast(dtype=torch.float32):
      # 如果有颜色条件和条件掩码，处理条件信息
      if color_cond is not None and cond_mask is not None:
        # 将条件颜色应用到输入，确保条件被应用到正确的位置
        cond_tokens = x.clone()
        # 在条件掩码位置应用颜色条件
        cond_positions = (cond_mask == 1).reshape(cond_mask.shape[0], -1)
        if cond_positions.any():
          cond_tokens[cond_positions] = color_cond[cond_positions]
        # 将条件作为额外输入传递给backbone
        logits = self.backbone(x, sigma, f_gray, cond_tokens)
      else:
        # 原始的无条件推理
        cond_tokens = None
        logits = self.backbone(x, sigma, f_gray)
      
    if (self.global_step + 1) % 100 == 0:
      # Save original colored image (x0)
      try:
        import torchvision
        import os

        # 只在主进程(rank 0)上保存可视化结果
        if self.trainer.is_global_zero:
          os.makedirs("visualization_outputs", exist_ok=True)

          # Convert from indices to images using VQGAN
          colored_img = self.chroma_vqgan.decode_img(x0, f_gray)  # Decode the original colored image
          # Create a modified version of x where mask tokens (16384) are replaced with 0
          x_modified = x.clone()
          x_modified[x_modified == self.mask_index] = 0
          noisy_img = self.chroma_vqgan.decode_img(x_modified, f_gray)  # Decode the noisy image

          # Convert model output logits to indices and then to image
          if self.parameterization != 'sedd':
            output_indices = logits.argmax(dim=-1)
            output_img = self.chroma_vqgan.decode_img(output_indices, f_gray)

            # 归一化图像到0-1范围
            colored_img_norm = torch.clamp((colored_img + 1) / 2, 0, 1)  # Assuming images are in [-1,1] range
            noisy_img_norm = torch.clamp((noisy_img + 1) / 2, 0, 1)
            output_img_norm = torch.clamp((output_img + 1) / 2, 0, 1)

            # 如果有颜色条件，也将其可视化
            if color_cond is not None and cond_mask is not None and cond_mask.any():
              cond_img = self.chroma_vqgan.decode_img(color_cond, f_gray)
              cond_img_norm = torch.clamp((cond_img + 1) / 2, 0, 1)
              
              # 将条件tokens也可视化（即带条件的输入tokens）
              if cond_tokens is not None:
                cond_tokens_modified = cond_tokens.clone()
                cond_tokens_modified[cond_tokens_modified == self.mask_index] = 0
                cond_tokens_img = self.chroma_vqgan.decode_img(cond_tokens_modified, f_gray)
                cond_tokens_img_norm = torch.clamp((cond_tokens_img + 1) / 2, 0, 1)
              color_condition_save = (color_condition + 1) / 2
              # # 创建保存路径
              # save_dir = f"visualization_outputs/step_{self.global_step}_cond"
              # os.makedirs(save_dir, exist_ok=True)
              
              # # 保存条件图像
              # for i in range(min(16, x0.shape[0])):
              #   # 组合原图、灰度图、条件图、条件tokens图和输出图
              #   if cond_tokens is not None:
              #     combined = torch.cat([colored_img_norm[i], noisy_img_norm[i], 
              #                          color_condition_save[i], output_img_norm[i]], dim=2)  # 水平连接
              #   else:
              #     combined = torch.cat([colored_img_norm[i], noisy_img_norm[i], 
              #                          cond_img_norm[i], output_img_norm[i]], dim=2)  # 水平连接
              #   torchvision.utils.save_image(combined, f"{save_dir}/img_{i}_combined.png")
              self._log_visualization_to_wandb(
                colored_img_norm, noisy_img_norm, output_img_norm,
                color_cond, cond_mask, cond_tokens, color_condition_save, f_gray
            )
            else:
              self._log_visualization_to_wandb(
                colored_img_norm, noisy_img_norm, output_img_norm,
                color_cond, cond_mask, cond_tokens, color_condition_save, f_gray
            )
              # # 创建保存路径
              # save_dir = f"visualization_outputs/step_{self.global_step}"
              # os.makedirs(save_dir, exist_ok=True)
              
              # # 创建组合图像（每个样本的三个版本水平排列）
              # for i in range(min(16, x0.shape[0])):
              #   combined = torch.cat([colored_img_norm[i], noisy_img_norm[i], output_img_norm[i]], dim=2)  # 水平连接
              #   torchvision.utils.save_image(combined, f"{save_dir}/img_{i}_combined.png")

            print(f"Saved batch visualization grid with {x0.shape[0]} images at step {self.global_step}")
            
            # 上传可视化到WandB
            
      except Exception as e:
        print(f"Error saving visualization: {e}")
    if self.parameterization == 'subs':
      return self._subs_parameterization(logits=logits,
                                         xt=x,
                                         cond_mask=cond_mask,
                                         color_cond=color_cond)
    elif self.parameterization == 'sedd':
      return self._sedd_parameterization(logits=logits,
                                         xt=x,
                                         sigma=sigma,
                                         cond_mask=cond_mask,
                                         color_cond=color_cond)
    elif self.parameterization == 'd3pm':
      return self._d3pm_parameterization(logits=logits,
                                         cond_mask=cond_mask,
                                         color_cond=color_cond)
    return logits

  def _d3pm_loss(self, model_output, xt, x0, t, cond_mask=None, color_cond=None):
    dt = 1 / self.T

    if torch.is_tensor(t):
      t = t[:, None]
      assert t.ndim == 2
      t = t.clamp(0., 1. - 1e-4)
    alpha_t = 1 - t + torch.zeros_like(xt)
    alpha_s = 1 - (t - dt) + torch.zeros_like(xt)

    log_x_theta_at_x0 = torch.gather(
      model_output, -1, x0[:, :, None]).squeeze(-1)
    log_x_theta_at_m = model_output[:, :, self.mask_index]
    x_theta_at_m = log_x_theta_at_m.exp()
    
    term_1_coef = dt / t
    term_1_log_nr = torch.log(alpha_t * x_theta_at_m / t + 1)
    term_1_log_dr = log_x_theta_at_x0
    
    term_2_coef = 1 - dt / t
    term_2_log_nr = term_1_log_nr
    term_2_log_dr = torch.log(alpha_s * x_theta_at_m / (t - dt) + 1)

    L_vb_masked = (
      term_1_coef * (term_1_log_nr - term_1_log_dr)
      + term_2_coef * (term_2_log_nr - term_2_log_dr))

    # 使用掩码选择哪些位置计算损失
    mask_positions = (xt == self.mask_index)
    
    # 如果有条件掩码，确保条件区域也计算损失
    if cond_mask is not None and color_cond is not None:
      cond_positions = (cond_mask == 1)
      if torch.any(cond_positions):
        # 计算条件区域的损失，可以使用加权方式
        cond_weight = 1.0  # 可以调整这个权重
        
        # 创建条件区域损失掩码：保留原来的mask位置，并加上条件位置（可能有重叠）
        combined_mask = mask_positions | cond_positions
        
        # 创建损失权重矩阵
        loss_weight = torch.ones_like(xt, dtype=L_vb_masked.dtype)
        # 对条件区域增加权重
        loss_weight[cond_positions] = cond_weight
        
        # 应用组合掩码和权重
        L_vb = L_vb_masked * combined_mask * loss_weight
    else:
      # 原始处理
      L_vb = L_vb_masked * mask_positions

    return self.T * L_vb

  def _compute_loss(self, batch, prefix):
    if 'attention_mask' in batch:
      attention_mask = batch['attention_mask']
    else:
      attention_mask = None
    losses = self._loss(batch, attention_mask)
    loss = losses.loss

    if prefix == 'train':
      self.train_metrics.update(losses.nlls, losses.token_mask)
      metrics = self.train_metrics
    elif prefix == 'val':
      self.valid_metrics.update(losses.nlls, losses.token_mask)
      metrics = self.valid_metrics
    elif prefix == 'test':
      self.test_metrics.update(losses.nlls, losses.token_mask)
      metrics = self.test_metrics
    else:
      raise ValueError(f'Invalid prefix: {prefix}')

    self.log_dict(metrics,
                  on_step=False,
                  on_epoch=True,
                  sync_dist=True)
    return loss

  def on_train_epoch_start(self):
    self.backbone.train()
    self.noise.train()

  def training_step(self, batch, batch_idx):
    loss = self._compute_loss(batch, prefix='train')
    self.log(name='trainer/loss',
             value=loss.item(),
             on_step=True,
             on_epoch=False,
             sync_dist=True)
    return loss

  def on_validation_epoch_start(self):
    if self.ema:
      self.ema.store(itertools.chain(
        self.backbone.parameters(),
        self.noise.parameters()))
      self.ema.copy_to(itertools.chain(
        self.backbone.parameters(),
        self.noise.parameters()))
    self.backbone.eval()
    self.noise.eval()
    assert self.valid_metrics.nll.mean_value == 0
    assert self.valid_metrics.nll.weight == 0

  def validation_step(self, batch, batch_idx):
    return self._compute_loss(batch, prefix='val')

  def on_validation_epoch_end(self):
    if ((self.config.eval.compute_perplexity_on_sanity
         or not self.trainer.sanity_checking)
         and self.config.eval.generate_samples
         and not self.parameterization == 'ar'):
      # TODO(justin): implement sampling and kv cache for AR
      samples, text_samples = None, None
      for _ in range(
        self.config.sampling.num_sample_batches):
        samples = self._sample()
        # Decode the samples to be re-tokenized by eval model
        text_samples = self.tokenizer.batch_decode(samples)
        if self.config.eval.compute_generative_perplexity:
          self.compute_generative_perplexity(text_samples)
      if self.trainer.global_rank == 0 and hasattr(
        self.trainer.logger, 'log_table'):
        # Log the last generated samples
        text_samples = text_samples[
          : self.config.sampling.num_sample_log]
        self.trainer.logger.log_table(
          key=f'samples@global_step{self.global_step}',
          columns=['Generated Samples'],
          data=[[s] for s in text_samples])
      if self.config.eval.compute_generative_perplexity:
        self.log('val/gen_ppl',
                 self.gen_ppl_metric,
                 on_epoch=True,
                 on_step=False,
                 sync_dist=True)
    if self.ema:
      self.ema.restore(
        itertools.chain(self.backbone.parameters(),
                        self.noise.parameters()))

  def configure_optimizers(self):
    # TODO(yair): Lightning currently giving this warning when using `fp16`:
    #  "Detected call of `lr_scheduler.step()` before `optimizer.step()`. "
    #  Not clear if this is a problem or not.
    #  See: https://github.com/Lightning-AI/pytorch-lightning/issues/5558
    optimizer = torch.optim.AdamW(
      itertools.chain(self.backbone.parameters(),
                      self.noise.parameters()),
      lr=self.config.optim.lr,
      betas=(self.config.optim.beta1,
             self.config.optim.beta2),
      eps=self.config.optim.eps,
      weight_decay=self.config.optim.weight_decay)

    scheduler = hydra.utils.instantiate(
      self.config.lr_scheduler, optimizer=optimizer)
    scheduler_dict = {
      'scheduler': scheduler,
      'interval': 'step',
      'monitor': 'val/loss',
      'name': 'trainer/lr',
    }
    return [optimizer], [scheduler_dict]

  @torch.no_grad()
  def eval_retokenize(self, text_samples, max_length):
    """Retokenizes samples for the eval model.
    
    Args:
        text_samples: List of sentences generated by the model.
    Returns:
        samples: Samples re-tokenized for the eval model
        attn_mask: Attention mask for the eval model
        eval_context_size: Size of the context for the eval model
    """
    if 'llama2' in self.gen_ppl_eval_model_name_or_path:
      tokenizer_kwargs = {
        'text_samples': text_samples,
        'return_tensors': 'pt',
        'return_token_type_ids': False,
        'return_attention_mask': True,
        'truncation': True,
        'padding': True,
        'max_length': max_length,
      }
      eval_context_size = 4096
    else:
      tokenizer_kwargs = {
        'return_tensors': 'pt',
        'return_token_type_ids': False,
        'return_attention_mask': True,
        'truncation': True,
        'padding': True,
        'max_length': max_length,
      }
      eval_context_size = 1024
    samples = self.eval_model_tokenizer(
      text_samples, ** tokenizer_kwargs)
    attn_mask = samples['attention_mask']
    samples = samples['input_ids']
    if 'llama2' not in self.gen_ppl_eval_model_name_or_path:
      attn_mask = attn_mask.to(self.device)
      samples = samples.to(self.device)      
    return samples, attn_mask, eval_context_size

  @torch.no_grad()
  def compute_generative_perplexity(
    self,
    text_samples: typing.List[str],
    retokenize: bool = True,
    max_length: typing.Optional[int] = None) -> None:
    """Compute the generative perplexity of the model.

    Args:
        text_samples: List of sentences generated by the model.
    
    Returns:
        Perplexity of the generated text under a different
        pre-trained AR model (e.g., GPT2).
    """
    os.environ['TOKENIZERS_PARALLELISM'] = 'false'
    eval_model = transformers.AutoModelForCausalLM.from_pretrained(
      self.gen_ppl_eval_model_name_or_path).eval()
    if max_length is None:
      max_length = self.config.model.length
    if 'llama2' not in self.gen_ppl_eval_model_name_or_path:
      eval_model = eval_model.to(self.device)
    # Re-tokenize using eval model's tokenizer
    if retokenize:
      (samples, attn_mask,
       eval_context_size) = self.eval_retokenize(
         text_samples, max_length=max_length)
    else:
      samples = text_samples
      attn_mask = torch.ones(samples.shape).to(self.device)
      eval_context_size = samples.shape[-1]
    batch_size = min(
      self.config.eval.perplexity_batch_size,
      samples.shape[0])
    num_batches = samples.shape[0] // batch_size
    for i in range(num_batches):
      _samples = torch.split(
        samples[i * batch_size: (i + 1) * batch_size],
        eval_context_size,
        dim=-1)
      _attn_mask = torch.split(
        attn_mask[i * batch_size: (i + 1) * batch_size],
        eval_context_size,
        dim=-1)
      for (sample_chunk, attn_mask_chunk) in zip(
        _samples, _attn_mask):
        logits = eval_model(
          sample_chunk, attention_mask=attn_mask_chunk)[0]
        logits = logits.transpose(-1, -2)
        
        nlls = F.cross_entropy(logits[..., :-1],
                               sample_chunk[..., 1:],
                               reduction='none')
        first_eos = (sample_chunk == self.eval_model_tokenizer\
                     .eos_token_id).cumsum(-1) == 1
        token_mask = (
          sample_chunk
          != self.eval_model_tokenizer.eos_token_id)
        self.gen_ppl_metric.update(
          nlls, first_eos[..., 1:] + token_mask[..., 1:])

  def q_xt(self, x, move_chance):
    """Computes the noisy sample xt.

    Args:
      x: int torch.Tensor with shape (batch_size,
          diffusion_model_input_length), input. 
      move_chance: float torch.Tensor with shape (batch_size, 1).
    """
    move_indices = torch.rand(
      * x.shape, device=x.device) < move_chance
    xt = torch.where(move_indices, self.mask_index, x)
    return xt

  def _sample_prior(self, *batch_dims):
    return self.mask_index * torch.ones(
      * batch_dims, dtype=torch.int64)

  def _ddpm_caching_update(self, x, t, dt, p_x0=None, color_cond=None, cond_mask=None, f_gray=None):
    assert self.config.noise.type == 'loglinear'
    sigma_t, _ = self.noise(t)
    if t.ndim > 1:
      t = t.squeeze(-1)
    assert t.ndim == 1
    move_chance_t = t[:, None, None]
    move_chance_s = (t - dt)[:, None, None]
    assert move_chance_t.ndim == 3, move_chance_t.shape
    if p_x0 is None:
      # 传递条件参数给forward函数
      p_x0 = self.forward(x, sigma_t, f_gray=f_gray, color_cond=color_cond, cond_mask=cond_mask).exp()
    
    assert move_chance_t.ndim == p_x0.ndim
    q_xs = p_x0 * (move_chance_t - move_chance_s)
    q_xs[:, :, self.mask_index] = move_chance_s[:, :, 0]
    _x = _sample_categorical(q_xs)
    
    copy_flag = (x != self.mask_index).to(x.dtype)
    return p_x0, copy_flag * x + (1 - copy_flag) * _x

  def _ddpm_update(self, x, t, dt, color_cond=None, cond_mask=None, f_gray=None):
    sigma_t, _ = self.noise(t)
    sigma_s, _ = self.noise(t - dt)
    if sigma_t.ndim > 1:
      sigma_t = sigma_t.squeeze(-1)
    if sigma_s.ndim > 1:
      sigma_s = sigma_s.squeeze(-1)
    assert sigma_t.ndim == 1, sigma_t.shape
    assert sigma_s.ndim == 1, sigma_s.shape
    move_chance_t = 1 - torch.exp(-sigma_t)
    move_chance_s = 1 - torch.exp(-sigma_s)
    move_chance_t = move_chance_t[:, None, None]
    move_chance_s = move_chance_s[:, None, None]
    unet_conditioning = sigma_t
    # 传递条件参数给forward函数
    log_p_x0 = self.forward(x, unet_conditioning, f_gray=f_gray, color_cond=color_cond, cond_mask=cond_mask)
    assert move_chance_t.ndim == log_p_x0.ndim
    # Technically, this isn't q_xs since there's a division
    # term that is missing. This division term doesn't affect
    # the samples.
    q_xs = log_p_x0.exp() * (move_chance_t
                             - move_chance_s)
    q_xs[:, :, self.mask_index] = move_chance_s[:, :, 0]
    _x = _sample_categorical(q_xs)

    copy_flag = (x != self.mask_index).to(x.dtype)
    return copy_flag * x + (1 - copy_flag) * _x

  def _ar_sampler(self, bsz):
    # precompute token buffer
    num_pred_tokens = self.config.model.length - 1
    x = torch.zeros(
      (bsz, num_pred_tokens + 1),
      dtype=torch.long,
      device=self.device)
    x[:, 0] = self.tokenizer.bos_token_id
    # precompute noise
    noise = (torch.distributions.Gumbel(0, 1)
             .sample((bsz, num_pred_tokens, self.vocab_size))
             .to(self.device))
    for i in range(num_pred_tokens):
      next_logits = self.forward(x[:, :i + 1], None)[:, -1]
      y = (next_logits + noise[:, i]).argmax(-1)
      x[:, i + 1] = y
    return x

  @torch.no_grad()
  def _sample(self, num_steps=None, eps=1e-5, color_cond=None, cond_mask=None, f_gray=None):
    """Generate samples from the model with optional color condition guidance.
    
    Args:
      num_steps: Number of diffusion steps
      eps: Small value to avoid division by zero
      color_cond: Optional color condition tokens, shape [batch_size, length]
      cond_mask: Optional mask for color conditions, shape [batch_size, length]
      f_gray: Grayscale image features
      
    Returns:
      Generated color tokens
    """
    # Lightning auto-casting is not working in this method for some reason
    if self.parameterization == 'ar':
      return self._ar_sampler(batch_size_per_gpu)
    
    batch_size_per_gpu = self.config.loader.eval_batch_size
    
    # Initialize with mask tokens
    x = self._sample_prior(
      batch_size_per_gpu,
      self.config.model.length).to(self.device)
      
    # 如果提供了颜色条件和掩码，在采样开始前应用条件
    if color_cond is not None and cond_mask is not None:
      # 将指定条件位置的mask token替换为颜色条件
      cond_positions = (cond_mask == 1)
      if cond_positions.any():
        x[cond_positions] = color_cond[cond_positions]
    
    if num_steps is None:
      num_steps = self.config.sampling.steps
      
    timesteps = torch.linspace(
      1, eps, num_steps + 1, device=self.device)
    dt = (1 - eps) / num_steps
    p_x0_cache = None

    for i in range(num_steps):
      t = timesteps[i] * torch.ones(
        x.shape[0], 1, device=self.device)
      if self.sampler == 'ddpm':
        x = self._ddpm_update(x, t, dt, color_cond, cond_mask)
      elif self.sampler == 'ddpm_cache':
        p_x0_cache, x_next = self._ddpm_caching_update(
          x, t, dt, p_x0=p_x0_cache, color_cond=color_cond, cond_mask=cond_mask)
        if (not torch.allclose(x_next, x)
            or self.time_conditioning):
          # Disable caching
          p_x0_cache = None
        x = x_next
      else:
        x = self._analytic_update(x, t, dt, color_cond, cond_mask)
        
      # 确保条件位置的颜色始终保持不变
      if color_cond is not None and cond_mask is not None:
        cond_positions = (cond_mask == 1)
        if cond_positions.any():
          x[cond_positions] = color_cond[cond_positions]

    if self.config.sampling.noise_removal:
      t = timesteps[-1] * torch.ones(x.shape[0], 1,
                                     device=self.device)
      if self.sampler == 'analytic':
        x = self._denoiser_update(x, t, color_cond, cond_mask)
      else:
        unet_conditioning = self.noise(t)[0]
        x = self.forward(x, unet_conditioning, f_gray=f_gray, color_cond=color_cond, cond_mask=cond_mask).argmax(dim=-1)
        
      # 最后一次确保条件位置的颜色保持不变
      if color_cond is not None and cond_mask is not None:
        cond_positions = (cond_mask == 1)
        if cond_positions.any():
          x[cond_positions] = color_cond[cond_positions]
          
    return x

  def restore_model_and_sample(self, num_steps=None, eps=1e-5, color_cond=None, cond_mask=None, f_gray=None):
    """Generate samples from the model with optional color conditions."""
    # Lightning auto-casting is not working in this method for some reason
    if self.ema:
      self.ema.store(itertools.chain(
        self.backbone.parameters(),
        self.noise.parameters()))
      self.ema.copy_to(itertools.chain(
        self.backbone.parameters(),
        self.noise.parameters()))
    self.backbone.eval()
    self.noise.eval()
    samples = self._sample(num_steps=num_steps, eps=eps, 
                          color_cond=color_cond, cond_mask=cond_mask, f_gray=f_gray)
    if self.ema:
      self.ema.restore(itertools.chain(
        self.backbone.parameters(),
        self.noise.parameters()))
    self.backbone.train()
    self.noise.train()
    return samples

  def get_score(self, x, sigma, f_gray=None, color_cond=None, cond_mask=None):
    model_output = self.forward(x, sigma, f_gray, color_cond=color_cond, cond_mask=cond_mask)
    if self.parameterization == 'subs':
      # score(x, t) = p_t(y) / p_t(x)
      # => log score(x, t) = log p_t(y) - log p_t(x)
      
      # case 1: x = masked
      #   (i) y = unmasked
      #     log score(x, t) = log p_\theta(x)|_y + log k
      #     where k = exp(- sigma) / (1 - exp(- sigma))
      #   (ii) y = masked
      #     log score(x, t) = 0

      # case 2: x = unmasked
      #   (i) y != masked, y != x
      #     log score(x_i, t) = - inf
      #   (ii) y = x 
      #     log score(x_i, t) = 0
      #   (iii) y = masked token
      #     log score(x_i, t) = - log k
      #     where k = exp(- sigma) / (1 - exp(- sigma))
      
      log_k = - torch.log(torch.expm1(sigma)).squeeze(-1)
      assert log_k.ndim == 1
      
      masked_score = model_output + log_k[:, None, None]
      masked_score[:, :, self.mask_index] = 0

      unmasked_score = self.neg_infinity * torch.ones_like(
        model_output)
      unmasked_score = torch.scatter(
        unmasked_score,
        -1,
        x[..., None],
        torch.zeros_like(unmasked_score[..., :1]))
      unmasked_score[:, :, self.mask_index] = - (
        log_k[:, None] * torch.ones_like(x))
      
      masked_indices = (x == self.mask_index).to(
        model_output.dtype)[:, :, None]
      model_output = (
        masked_score * masked_indices
        + unmasked_score * (1 - masked_indices))
      
      # 如果有颜色条件和掩码，保持条件位置不变
      if color_cond is not None and cond_mask is not None:
        cond_positions = (cond_mask == 1)
        if cond_positions.any():
          # 在条件位置，将得分集中在条件颜色上
          for b in range(model_output.shape[0]):
            for pos in torch.nonzero(cond_positions[b]):
              pos_idx = pos.item()
              # 将除了条件颜色以外的所有得分设置为-无穷大
              model_output[b, pos_idx] = self.neg_infinity
              if pos_idx < color_cond.shape[1]:  # 确保pos_idx在color_cond范围内
                color_idx = color_cond[b, pos_idx].item()
                model_output[b, pos_idx, color_idx] = 0  # 只保留条件颜色的得分
          
    return model_output.exp()

  def _staggered_score(self, score, dsigma):
    score = score.clone()
    extra_const = (1 - dsigma.exp()) * score.sum(dim=-1)
    score *= dsigma.exp()[:, None]
    score[..., self.mask_index] += extra_const
    return score

  def _analytic_update(self, x, t, step_size, color_cond=None, cond_mask=None):
    curr_sigma, _ = self.noise(t)
    next_sigma, _ = self.noise(t - step_size)
    dsigma = curr_sigma - next_sigma
    score = self.get_score(x, curr_sigma, color_cond=color_cond, cond_mask=cond_mask)
    stag_score = self._staggered_score(score, dsigma)
    probs = stag_score * self._transp_transition(x, dsigma)
    return _sample_categorical(probs)

  def _denoiser_update(self, x, t, color_cond=None, cond_mask=None):
    sigma, _ = self.noise(t)
    score = self.get_score(x, sigma, color_cond=color_cond, cond_mask=cond_mask)
    stag_score = self._staggered_score(score, sigma)
    probs = stag_score * self._transp_transition(x, sigma)
    probs[..., self.mask_index] = 0
    samples = _sample_categorical(probs)
    return samples

  def _transp_transition(self, i, sigma):
    sigma = _unsqueeze(sigma, reference=i[..., None])
    edge = torch.exp(-sigma) * F.one_hot(
      i, num_classes=self.vocab_size)
    edge += torch.where(i == self.mask_index,
                        1 - torch.exp(-sigma).squeeze(-1),
                        0)[..., None]
    return edge

  def _sample_t(self, n, device):
    _eps_t = torch.rand(n, device=device)
    if self.antithetic_sampling:
      offset = torch.arange(n, device=device) / n
      _eps_t = (_eps_t / n + offset) % 1
    t = (1 - self.sampling_eps) * _eps_t + self.sampling_eps
    if self.importance_sampling:
      return self.noise.importance_sampling_transformation(t)
    return t

  def _maybe_sub_sample(self, x0, attention_mask):
    seqlen = x0.shape[1]
    if seqlen > self.config.model.length:
      assert seqlen == 2 * self.config.model.length
      # cropping is needed for text8-crop dataset
      # try the same starting point for now
      start = np.random.choice(self.config.model.length)
      end = start + self.config.model.length
      input_tokens = x0[:, start: end]
      output_tokens = x0[:, start + 1: end + 1]
      new_attention_mask = attention_mask[:, start: end]

      # Helps with validation PPL, since the val
      # examples will all start and end with BOS/EOS
      input_tokens[:, 0] = self.tokenizer.bos_token_id
      output_tokens[:, -1] = self.tokenizer.eos_token_id
    elif self.parameterization == 'ar':
      input_tokens = x0[:, :-1]
      output_tokens = x0[:, 1:]
      new_attention_mask = attention_mask[:, 1:]
    else:
      input_tokens = x0
      output_tokens = None
      new_attention_mask = attention_mask
    return input_tokens, output_tokens, new_attention_mask

  def _reconstruction_loss(self, x0):
    t0 = torch.zeros(x0.shape[0], dtype=self.dtype,
                     device=self.device)
    assert self.config.noise.type == 'loglinear'
    # The above assert is for d3pm parameterization
    unet_conditioning = self.noise(t0)[0][:, None]
    model_output_t0 = self.forward(x0, unet_conditioning)
    return - torch.gather(input=model_output_t0,
                          dim=-1,
                          index=x0[:, :, None]).squeeze(-1)

  def _forward_pass_diffusion(self, x0, f_gray=None, color_cond=None, cond_mask=None, color_condition=None):
    t = self._sample_t(x0.shape[0], x0.device)
    if self.T > 0:
      t = (t * self.T).to(torch.int)
      t = t / self.T
      # t \in {1/T, 2/T, ..., 1}
      t += (1 / self.T)

    if self.change_of_variables: #false
      unet_conditioning = t[:, None]
      f_T = torch.log1p(- torch.exp(- self.noise.sigma_max))
      f_0 = torch.log1p(- torch.exp(- self.noise.sigma_min))
      move_chance = torch.exp(f_0 + t * (f_T - f_0))
      move_chance = move_chance[:, None]
    else:
      sigma, dsigma = self.noise(t)
      unet_conditioning = sigma[:, None]
      move_chance = 1 - torch.exp(-sigma[:, None])

    xt = self.q_xt(x0, move_chance) # 按给定概率把输入序列里的部分 token 替换成“mask token”
    model_output = self.forward(xt, 
                                unet_conditioning, 
                                f_gray=f_gray, 
                                x0=x0, 
                                color_cond=color_cond, 
                                cond_mask=cond_mask, 
                                color_condition=color_condition)
    utils.print_nans(model_output, 'model_output')

    if self.parameterization == 'sedd':
      return dsigma[:, None] * self._score_entropy(
        model_output, sigma[:, None], xt, x0)
    
    if self.T > 0:
      diffusion_loss = self._d3pm_loss(
        model_output=model_output, xt=xt, x0=x0, t=t, cond_mask=cond_mask, color_cond=color_cond)
      if self.parameterization == 'd3pm':
        reconstruction_loss = self._reconstruction_loss(x0)
      elif self.parameterization == 'subs':
        reconstruction_loss = 0
      return reconstruction_loss + diffusion_loss
    
    # SUBS parameterization, continuous time.
    log_p_theta = torch.gather(
      input=model_output,
      dim=-1,
      index=x0[:, :, None]).squeeze(-1)
    
    if self.change_of_variables or self.importance_sampling:
      return log_p_theta * torch.log1p(
        - torch.exp(- self.noise.sigma_min))
    
    return - log_p_theta * (
      dsigma / torch.expm1(sigma))[:, None]

  def _loss(self, batch, attention_mask):
    # (input_tokens, output_tokens,
    #  attention_mask) = self._maybe_sub_sample(
    #    x0, attention_mask)
    x_gray = batch['grayscale']
    x_color = batch['original']
    color_idx, f_gray = self.chroma_vqgan.encode(x_color=x_color, x_gray=x_gray)
    
    # 处理条件上色信息
    color_cond = None
    cond_mask = None
    if 'color_condition' in batch and 'condition_mask' in batch:
      # 获取颜色条件和掩码
      color_condition = batch['color_condition']
      condition_mask = batch['condition_mask']
      
      if torch.any(condition_mask > 0):
        # 对颜色条件进行编码
        color_cond, _ = self.chroma_vqgan.encode(x_color=color_condition, x_gray=x_gray)
        cond_mask = condition_mask
    else:
      assert 1==2
    if self.parameterization == 'ar':
      pass
      # logprobs = self.backbone(input_tokens, None)
      # loss = - logprobs.gather(
      #   -1, output_tokens[:, :, None])[:, :, 0]
    else:
      # 传递条件信息到扩散模型
      loss = self._forward_pass_diffusion(x0=color_idx, f_gray=f_gray, 
                                          color_cond=color_cond, cond_mask=cond_mask, color_condition=color_condition)
    
    # nlls = loss * attention_mask
    # count = attention_mask.sum()

    token_nll = loss.mean()
    # token_nll = batch_nll / count

    return Loss(loss=token_nll,
                nlls=token_nll,
                token_mask=attention_mask)

  def _score_entropy(self, log_score, sigma, xt, x0):
    """Computes the SEDD loss.

    Args:
      log_score: float torch.Tensor with shape (batch_size,
          diffusion_model_input_length, vocab_size),
          log score, output of the denoising network.
      xt: int torch.Tensor with shape (batch_size,
          diffusion_model_input_length), input.
      x0: int torch.Tensor with shape (batch_size,
          diffusion_model_input_length), input.
      sigma: float torch.Tensor with shape (batch_size, 1).

    Returns:
      loss with shape (batch_size, diffusion_model_input_length)
    """
    masked_indices = xt == self.mask_index

    expsig_minus_1 = torch.expm1(sigma).expand_as(xt)
    q_ratio = 1 / expsig_minus_1[masked_indices]

    words_that_were_masked = x0[masked_indices]

    neg_term = q_ratio * torch.gather(
      log_score[masked_indices],
      -1,
      words_that_were_masked[..., None]).squeeze(-1)
    score = log_score[masked_indices].exp()
    if self.mask_index == self.vocab_size - 1:
      pos_term = score[:, :-1].sum(dim=-1)
    else:
      pos_term = score[:, : self.mask_index].sum(
        dim=-1) + score[:, self.mask_index + 1:].sum(dim=-1)
    const = q_ratio * (q_ratio.log() - 1)

    entropy = torch.zeros(* xt.shape, device=xt.device)
    entropy[masked_indices] += pos_term - neg_term + const
    return entropy

  @torch.no_grad
  def sample_subs_guidance(
    self, n_samples, stride_length, num_strides, dt=0.001):
    ones = torch.ones(n_samples, dtype=self.dtype,
                      device=self.device)

    num_steps = int(1 / dt)
    sampling_steps = 0
    intermediate_tokens = []
    target = None
    for _ in range(num_strides + 1):
      p_x0_cache = None
      x = self._sample_prior(
        n_samples,
        self.config.model.length).to(self.device)
      if target is not None:
        x[:, : -stride_length] = target
      for i in range(num_steps + 1):
        p_x0_cache, x_next = self._ddpm_caching_update(
          x=x, t=(1 - i * dt) * ones, dt=dt, p_x0=p_x0_cache)
        if (not torch.allclose(x_next, x)
            or self.time_conditioning):
          p_x0_cache = None
          sampling_steps += 1
        x = x_next
      x = self.forward(x, 0 * ones).argmax(dim=-1)
      intermediate_tokens.append(
        x[:, :stride_length].cpu().numpy())
      target = x[:, stride_length:]
    
    intermediate_tokens.append(target.cpu().numpy())
    intermediate_text_samples = []
    sequence_lengths = ((
      np.concatenate(intermediate_tokens, axis=1)[:, 1:]
      == self.tokenizer.eos_token_id).cumsum(-1) == 0).sum(-1)
    for i in range(2, len(intermediate_tokens) + 1):
      intermediate_text_samples.append(
        self.tokenizer.batch_decode(
          np.concatenate(intermediate_tokens[:i], axis=1)))
    return (sampling_steps, intermediate_text_samples,
            sequence_lengths)

  def restore_model_and_semi_ar_sample(
      self, stride_length, num_strides, dt=0.001):
    """Generate samples from the model."""
    # Lightning auto-casting is not working in this method for some reason
    if self.ema:
      self.ema.store(itertools.chain(
        self.backbone.parameters(),
        self.noise.parameters()))
      self.ema.copy_to(itertools.chain(
        self.backbone.parameters(),
        self.noise.parameters()))
    self.backbone.eval()
    self.noise.eval()
    (sampling_steps, samples,
     sequence_lengths) = self.sample_subs_guidance(
      n_samples=self.config.loader.eval_batch_size,
      stride_length=stride_length,
      num_strides=num_strides, 
      dt=dt)
    if self.ema:
      self.ema.restore(itertools.chain(
        self.backbone.parameters(),
        self.noise.parameters()))
    self.backbone.train()
    self.noise.train()
    return sampling_steps, samples, sequence_lengths
